import torch
from transformers import BertForSequenceClassification, BertTokenizer  # 明确导入Bert相关类
import numpy as np

# 1. 标签映射（与训练时完全一致）
LABEL_MAP = {
    "开心": 0,
    "伤心": 1,
    "厌恶": 2,
    "惊讶": 3,
    "生气": 4,
    "关心": 5,
    "疑问": 6,
    "平静": 7
}
id2label = {v: k for k, v in LABEL_MAP.items()}
num_classes = len(LABEL_MAP)

# 2. 加载BertTokenizer和模型（使用BertTokenizer而非AutoTokenizer）
# 注意：模型需与训练时使用的BERT版本一致（如bert-base-chinese）
model_name_or_path = r"D:\AI\model\bert-base-chinese"  # 替换为你的模型路径（本地微调后的路径）
tokenizer = BertTokenizer.from_pretrained(model_name_or_path)  # 明确使用BertTokenizer
model = BertForSequenceClassification.from_pretrained(
    model_name_or_path,
    num_labels=num_classes,
    id2label=id2label,
    label2id=LABEL_MAP
)

# 3. 加载训练好的.pth权重文件
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 加载权重并适配名称
state_dict = torch.load(r"params\1bert.pth", map_location=device)

# 创建新的状态字典，将fc.weight和fc.bias重命名为classifier.weight和classifier.bias
new_state_dict = {}
for key, value in state_dict.items():
    if key == "fc.weight":
        new_state_dict["classifier.weight"] = value
    elif key == "fc.bias":
        new_state_dict["classifier.bias"] = value
    else:
        # 保留BERT模型的其他权重名称（如果存在）
        new_state_dict[key] = value

# 加载适配后的权重
model.load_state_dict(new_state_dict, strict=False)  # 使用strict=False忽略不匹配的键
model.to(device)

model.eval()  # 切换为评估模式


# 4. 预处理函数（基于BertTokenizer的中文处理）
def preprocess_text(text, tokenizer, max_seq_len=128):
    """
    使用BertTokenizer处理中文文本
    自动完成：分词（字级+子词）、转换为input_ids、添加padding和attention_mask
    """
    # BertTokenizer的encode_plus方法可直接生成模型所需的所有输入
    encoding = tokenizer.encode_plus(
        text=text,
        max_length=max_seq_len,  # 与训练时的最大序列长度保持一致
        padding="max_length",  # 填充至max_seq_len
        truncation=True,  # 超过长度则截断
        return_tensors="pt",  # 返回PyTorch张量
        return_attention_mask=True  # 返回attention_mask（区分真实文本和padding）
    )
    return encoding  # 包含input_ids、attention_mask等键


# 5. 预测函数（与BERT模型适配）
def predict_sentiment(text, model, tokenizer, device, max_seq_len=30):
    # 预处理文本
    encoding = preprocess_text(text, tokenizer, max_seq_len)
    input_ids = encoding["input_ids"].to(device)
    attention_mask = encoding["attention_mask"].to(device)

    # 关闭梯度计算，执行推理
    with torch.no_grad():
        outputs = model(
            input_ids=input_ids,
            attention_mask=attention_mask
        )
        logits = outputs.logits  # BERT分类头的输出
        probabilities = torch.softmax(logits, dim=1).squeeze().cpu().numpy()  # 转换为概率

    # 解析预测结果
    predicted_label_id = np.argmax(probabilities)
    predicted_label = id2label[predicted_label_id]

    return {
        "text": text,
        "predicted_label": predicted_label,
        "confidence": float(probabilities[predicted_label_id]),
        "all_probabilities": {id2label[i]: float(p) for i, p in enumerate(probabilities)}
    }


# 6. 测试口语化文本案例
if __name__ == "__main__":
    # 测试用例：覆盖各类口语情感表达
    test_cases = [
        "哦！我被选中了！",
        "天哪，这真是太惊喜了！",
        "唉，又失败了，真让人难过",
        "气死我了，怎么能这样对我！",
        "嗯，这件事就这样吧",
        "哇塞！我居然赢了！",
        "呜呜呜，我的东西丢了",
        "什么？你再说一遍？",
        "还好吧，不算太坏",
        "太棒啦！我们成功了！"
    ]

    # 批量预测并打印结果
    print("===== 情感分类结果 =====")
    for text in test_cases:
        result = predict_sentiment(text, model, tokenizer, device)
        print(f"文本: {result['text']}")
        print(f"预测情感: {result['predicted_label']} (置信度: {result['confidence']:.4f})")
        print("详细概率:", {k: f"{v:.4f}" for k, v in result['all_probabilities'].items()})
        print("---")
